use num_bigint::BigUint;
use num_traits::{identities::Zero, One};
use rand::Rng;
use std::ops::Shr;

fn mod_exp(base: u128, exp: u128, modulus: u128) -> u128 {
    let base = BigUint::from(base);
    let mut exp = BigUint::from(exp);
    let modulus = BigUint::from(modulus);
    let mut result = BigUint::one();
    let mut base = base % &modulus;

    while !exp.is_zero() {
        if &exp & BigUint::one() == BigUint::one() {
            result = (result * &base) % &modulus;
        }
        base = (&base * &base) % &modulus;
        exp = exp.shr(1u8);
    }
    result.try_into().expect("Failed to convert BigUint to u128")
}

fn decompose(n: u128) -> (u128, u128) {
    let mut d = n - 1;
    let mut r = 0;

    while d % 2 == 0 {
        d /= 2;
        r += 1;
    }
    (d, r)
}

fn miller_rabin_test(n: u128, a: u128, d: u128, r: u128) -> bool {
    let mut x = mod_exp(a, d, n);
    if x == 1 || x == n - 1 {
        return true;
    }

    for _ in 0..r - 1 {
        x = mod_exp(x, 2, n);
        if x == n - 1 {
            return true;
        }
    }
    false
}

fn is_prime(n: u128, k: u32) -> bool {
    if n <= 1 {
        return false;
    }
    if n <= 3 {
        return true;
    }
    if n % 2 == 0 {
        return false;
    }

    let (d, r) = decompose(n);

    let mut rng = rand::thread_rng();
    for _ in 0..k {
        let a = rng.gen_range(2..n - 1);
        if !miller_rabin_test(n, a, d, r) {
            return false;
        }
    }
    true
}


fn gcd(mut a: u128, mut b: u128) -> u128 {
    while a % b != 0 {
        let r = a % b;
        a = b;
        b = r;
    }
    b
}

fn gcd_bigint(a: &BigUint, b: &BigUint) -> BigUint {
    let mut a_new = a.clone();
    let mut b_new = b.clone();
    while a_new.clone() % &b_new != BigUint::zero() {
        let r = a_new % &b_new;
        a_new = b_new;
        b_new = r;
    }
    b_new
}

const BASIC_PRIMES: [u128; 4] = [2, 3, 5, 7];

// fn pollard_rho(n: u128) -> u128 {
//     for &p in BASIC_PRIMES.iter() {
//         if n % p == 0 {
//             return p;
//         }
//     }
//     let mut rng = rand::thread_rng();
//     loop {
//         let mut x = rng.gen_range(1..n);
//         let mut y = x;
//         let mut factor = 1;

//         while factor == 1 {
//             x = ((x as u128).pow(2) + 1) % n;
//             y = ((y as u128).pow(2) + 1) % n;
//             y = ((y as u128).pow(2) + 1) % n;
//             factor = gcd((x as i128 - y as i128).abs() as u128, n);
//         }

//         if factor != n {
//             return factor;
//         }
//     }
// }

fn pollard_rho_bigint(number: u128) -> u128 {
    for &p in BASIC_PRIMES.iter() {
        if number % p == 0 {
            return p;
        }
    }
    let n = BigUint::from(number);
    loop {
        let mut rng = rand::thread_rng();
        let x = rng.gen_range(1..number);
        let mut x = BigUint::from(x);
        let mut y = x.clone();
        let mut factor = BigUint::one();

        while factor == BigUint::one() {
            x = (&x * &x + BigUint::one()) % &n;
            y = (&(&y * &y) + BigUint::one()) % &n;
            y = (&(&y * &y) + BigUint::one()) % &n;
            let diff = if x > y { &x - &y } else { &y - &x };
            factor = gcd_bigint(&diff, &n);
        }

        if factor != n {
            return factor.try_into().expect("Failed to convert BigUint to u128");
        }
    }
}

fn largest_prime_factor(n: u128) -> u128 {
    if is_prime(n, 5) {
        return n;
    }
    let mut factor = pollard_rho_bigint(n);
    while !is_prime(factor, 5) {
        factor = pollard_rho_bigint(factor);
    }
    let remaining = n / factor;
    if remaining == 1 {
        factor
    } else {
        let remaining_factor = largest_prime_factor(remaining);
        if factor > remaining_factor {
            factor
        } else {
            remaining_factor
        }
    }
}


pub fn find_max_prime_factor(number: u128) -> u128 {
    largest_prime_factor(number)
}
